#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

import ast
import inspect
import re
import types
from ast import *

def pretty_dump(node, annotate_fields = True, include_attributes = False):
    """
    Modified from ast.dump to pretty print the result.
    
    Return a formatted pretty_dump of the tree in *node*.  This is mainly useful for
    debugging purposes.  The returned string will show the names and the values
    for fields.  This makes the code impossible to evaluate, so if evaluation is
    wanted *annotate_fields* must be set to False.  Attributes such as line
    numbers and column offsets are not dumped by default.  If this is wanted,
    *include_attributes* can be set to True.
    """
    indent = '   '
    def _format(node, level = 0):
        if isinstance(node, AST):
            fields = [(a, _format(b, level + 1)) for a, b in iter_fields(node)]
            rv = '%s(%s' % (node.__class__.__name__, ', '.join(
                ('%s=%s' % field for field in fields)
                if annotate_fields else
                (b for a, b in fields)
            ))
            if include_attributes and node._attributes:
                rv += fields and ', ' or ' '
                rv += ', '.join('%s=%s' % (a, _format(getattr(node, a), level + 1))
                                for a in node._attributes)
            return "\n" + level * indent + rv + ')'
        elif isinstance(node, list):
            return '[%s]' % ', '.join(_format(x, level + 1) for x in node)
        return repr(node)
    if not isinstance(node, AST):
        raise TypeError('expected AST, got %r' % node.__class__.__name__)
    return _format(node)


class MetaEnhance(object):
    '''
    Enhances objects, making their functions easier to trace.
    Consistent use of meta-programming. Beware.
    '''
    def __init__(self, namespace):
        self.trace = False
        self.outFunction = 'print'
        self.namespace = globals()
        self.namespace.update(namespace)

    def traceit(self, other, methodNames = None):
        if methodNames is None:
            for method in self.methodList(other):
                self.enhanceObject(other, method)
        else:
            for name in methodNames:
                method = getattr(other, name)
                self.enhanceObject(other, method)


    ####################
    ## Private methods
    ####################
    def enhanceObject(self, other, method):
        source = self.getCleanSource(method)
        if not source[0].startswith("@classmethod"):
            name = method.__name__
            source = self.enhanceSourceToPrintIt(source, name)
            func = self.enhanceFunctionToTraceAssignements(source, name)
            setattr(other, method.__name__, types.MethodType(func, other))

    def enhanceObjectToTraceAssignements(self, other, method):
        func = self.enhanceFunctionToTraceAssignements(self.getCleanSource(method), method.__name__)
        setattr(other, method.__name__, types.MethodType(func, other))

    def enhanceFunctionToTraceAssignements(self, sourceLines, functionName):
        node = ast.parse(self.compileStr(sourceLines))
        self.traceEnhanceAssign(sourceLines, node)
        newNode = fix_missing_locations(FunctionVisitor().visit(node))
        self.traceEnhanceAssign(sourceLines, node)
        exec self.compile(newNode) in self.namespace
        return self.namespace[functionName]

    def compile(self, nodeOrStr):
        return compile(nodeOrStr, "<MetaEnhance.traceit>", "exec")

    def enhanceObjectToPrintIt(self, other, method):
        func = self.enhanceFunctionToPrintIt(self.getCleanSource(method), method.__name__)
        setattr(other, method.__name__, types.MethodType(func, other))


    def enhanceSourceToPrintIt(self, sourceLines, functionName):
        lines = sourceLines
        newLines = [lines[0]]
        lineCount = 1
        appendNext = None
        for line in lines[self.commentCount(lines):]:
            data = line.strip()
            if appendNext:
                text, linenumber = appendNext
                tuple = (self.outFunction, functionName, linenumber, self.quote(text))
                output = "%s('%s %d:%s')" % tuple
                newLines.append(self.indentOf(line) + output + "\n")
                appendNext = None
            if data.startswith("else:") or data.startswith("elif") or data.startswith("except"):
                appendNext = line.rstrip(), lineCount
            else:
                newLines.append(self.metaPrintLine(line.rstrip(), lineCount, functionName))
            newLines.append(line)
            lineCount += 1
        if len(newLines) == 1:
            newLines.append("    pass\n")
        self.traceEnhance(newLines)
        return newLines

    def commentCount(self, lines):
        afterDef = lines[1].strip()
        count = 1
        if afterDef.startswith("'''") or afterDef.startswith('"""'):
            commentType = afterDef[0:len("'''")]
            lines[1] = afterDef[len("'''"):]
            while count < len(lines):
                if lines[count].strip().endswith(commentType):
                    return count + 1
                count += 1
        return count

    def enhanceFunctionToPrintIt(self, sourceLines, functionName):
        newLines = self.enhanceSourceToPrintIt(sourceLines, functionName)
        exec self.compile(self.compileStr(newLines)) in self.namespace
        return self.namespace[functionName]

    def compileStr(self, newLines):
        return "".join(newLines)

    def traceEnhanceAssign(self, newLines, node):
        if self.trace:
            source = self.compileStr(newLines)
            print "compilint\n", source
            print pretty_dump(node)

    def traceEnhance(self, newLines):
        if self.trace:
            source = self.compileStr(newLines)
            print "compilint\n", source
            print pretty_dump(ast.parse(source))

    def metaPrintLine(self, line, linenumber, functionName):
        tuple = (self.outFunction, functionName, linenumber, self.quote(line))
        output = "%s('%s %d:%s')" % tuple
        return self.indentOf(line) + output + "\n"


    def quote(self, str):
        return re.sub(r"'", "\\'", str)

    def getCleanSource(self, function):
        source = inspect.getsourcelines(function)[0]
        firstline = source[0]
        stripSize = self.indentSize(firstline)
        return [self.removeIndent(line, stripSize) for line in source]


    def indentSize(self, line):
        return len(line) - len(line.lstrip())

    def indentOf(self, line):
        return line[0:self.indentSize(line)]

    def removeIndent(self, line, indentSize):
        return line[indentSize:]

    def isMethod(self, value):
        return inspect.ismethod(value)

    def methodList(self, target):
        methods = (getattr(target, method) for method in dir(target))
        return (method for method in methods if self.isMethod(method))


class PrintDict(dict):
    '''
    Dictionary that prints everything it sets
    '''
    def __setitem__(self, key, value):
        super(PrintDict, self).__setitem__(key, value)
        print "%s = %r" % (key, value)

class FunctionVisitor(ast.NodeTransformer):

    def __init__(self):
        self.variables = set()
        self.methodName = None

    @classmethod
    def dictionary(self):
        return PrintDict()

    @classmethod
    def returnAndPrint(self, methodName, value):
        self.output("return value of %s = %s\n" % (methodName, str(value)))
        return value

    @classmethod
    def output(self, value):
        print value

    def visit_Assign(self, node):
        for n in node.targets:
            visitor = VariablesVisitor()
            visitor.visit(n)
            self.variables.update(visitor.variables)
        self.generic_visit(node)
        return node

    def visit_Name(self, node):
        if node.id == 'self' or node.id not in self.variables:
            return node
        return Subscript(
            value = Name(id = '__data__', ctx = Load()),
            slice = Index(value = Str(s = node.id)),
            ctx = node.ctx
        )

    def visit_Return(self, node):
        newvalue = self.visit(node.value)
        node.value = Call(func =
                  Attribute(value =
                     Name(id = self.classname(), ctx =
                        Load()), attr = 'returnAndPrint', ctx =
                     Load()), args = [Str (s = self.methodName), newvalue], keywords = [], starargs = None, kwargs = None)
        return node

    def classname(self):
        return self.__class__.__name__


    def visit_FunctionDef(self, node):
        self.methodName = node.name
        for n in node.body:
            self.visit(n)
        before = [Expr(value =
               Call(func =
                  Attribute(value =
                     Name(id = self.classname(), ctx =
                        Load()), attr = 'output', ctx =
                     Load()), args = [
                     Str(s = ">>>> calling " + self.methodName)], keywords = [], starargs = None, kwargs = None))
                  , Assign(targets = [
                  Name(id = '__data__', ctx =
                     Store())], value =
               Call(func =
                  Attribute(value =
                     Name(id = self.classname(), ctx =
                        Load()), attr = 'dictionary', ctx =
                     Load()), args = [], keywords = [], starargs = None, kwargs = None))]
        argsNames = [i.id for i in node.args.args if i.id != 'self']
        for arg in argsNames:
            before.append(Assign(targets = [
                  Subscript(value =
                     Name(id = '__data__', ctx =
                        Load()), slice =
                     Index(value =
                        Str(s = arg)), ctx =
                     Store())], value =
                     Name(id = arg, ctx =
                          Load())))
        node.body = before + node.body
        return node


class VariablesVisitor(ast.NodeVisitor):
    def __init__(self):
        self.variables = set()

    def visit_Name(self, node):
        self.variables.add(node.id)
        self.generic_visit(node)
